import os

import tushare as ts
from pymongo import MongoClient
import json
import stock.config as cfg

class StockCatchNew:
    def __init__(self, todb=False, tocsv=False):
        self.todb = todb
        self.tocsv = tocsv
        f = open(os.getcwd() + "/../stock/config.json")
        data = f.read()
        f.close()
        conf = json.loads(data)
        self.ts_pro = ts.pro_api(conf['token'])
        if todb:
            #mdburl='mongodb://' + conf['dbuser'] + ':' + conf['dbpass'] + '@' + conf['dburl']
            mdburl='mongodb://' + conf['dburl']
            self.client = MongoClient(mdburl)
            self.stockdb = self.client['stock']

    # 获取信息
    def getDailyToday(self, today):
        df = self.ts_pro.daily(trade_date=today)
        if self.todb:
            dbname = 'newdaily-'+today
            self.stockdb[dbname].remove()
            self.stockdb[dbname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/" + "today-" + today + ".csv"
            df.to_csv(filename)

    def get_daily_full(self, today):
        df = self.ts_pro.daily_basic(trade_date=today)
        if self.todb:
            dbname = 'daily_full-'+today
            self.stockdb[dbname].remove()
            self.stockdb['daily_full'].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/" + "daily_full-" + today + ".csv"
            df.to_csv(filename)

    def getDaily(self, code, start, end):
        df = self.ts_pro.daily(ts_code=code, start_date=start, end_date=end)
        if self.todb:
            cname = 'newdaily-'+code
            try:
                self.stockdb[cname].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb[cname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/newdaily-" + code + "-" + start + "-" + end + ".csv"
            df.to_csv(filename)

    def getProBar(self, code, start, end):
        df = self.ts_pro.pro_bar(ts_code=code, adj='qfq', start_date=start, end_date=end)
        if self.todb:
            cname = 'probar-' + code
            try:
                self.stockdb[cname].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb[cname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/probar-" + code + "-" + end + ".csv"
            df.to_csv(filename)

    def getDailyTrade(self, trade):
        df = self.ts_pro.daily(trade_date=trade)
        if self.todb:
            cname = 'newdaily-'+trade
            try:
                self.stockdb[cname].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb[cname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/" + 'newdaily' + "-" + trade + ".csv"
            df.to_csv(filename)

    def getIndexDaily(self, code, start, end):
        df = self.ts_pro.index_daily(ts_code=code, start_date=start, end_date=end)
        if self.todb:
            cname = 'newdaily-' + code
            try:
                self.stockdb[cname].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb[cname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/newdaily-" + code + "-" + start + "-" + end + ".csv"
            df.to_csv(filename)

    def getLimitList(self, today):
        df = self.ts_pro.limit_list()
        if self.todb:
            cname = 'limit_list-' + today
            self.stockdb[cname].remove()
            self.stockdb[cname].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            filename = cfg.t_localpath + "/" + 'limit_list_' + today + ".csv"
            df.to_csv(filename)

    # 所有股票列表
    def getBasic(self):
        df = self.ts_pro.stock_basic()
        if self.todb:
            try:
                print("get basic to database")
                self.stockdb['basic'].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb['basic'].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            df.to_csv(cfg.t_localpath + "/basic.csv")

    # 所有上市公司，信息
    def getCompany(self):
        df = self.ts_pro.stock_company()
        if self.todb:
            try:
                self.stockdb['company'].remove()
            except Exception as e:
                print("not exist , not not to remove it")
            self.stockdb['company'].insert_many(json.loads(df.to_json(orient='records')))
        if self.tocsv:
            df.to_csv(cfg.t_localpath + "/company.csv")
